import pathlib
from typing import ClassVar, override

import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm

from ..processors import GensimLDAStage
from ..utils import OutputManager
from .interface import Marshaler

__all__ = ["LDAVisualizationMarshaler", "LDAStatisticsMarshaler", "LDAKeywordsExcelMarshaler"]


class LDAVisualizationMarshaler(Marshaler[GensimLDAStage]):
    _output_dir: pathlib.Path

    def __init__(self, output_manager: OutputManager) -> None:
        self._output_dir = output_manager.filepath("lda-visualization")
        if not self._output_dir.exists():
            self._output_dir.mkdir(parents=True)

    @override
    def marshal(self, stage: GensimLDAStage) -> None:
        # These imports cannot be placed at the top of the file, because they will trigger a deprecation warning.
        #
        # I don't know why this is happening, but we better do so.
        import pyLDAvis

        for index, preset in tqdm(enumerate(stage.visualization_presets), desc="Visualizing LDA"):
            with open(self._output_dir / f"{preset.num_topics}_topics.html", "w") as f:
                pyLDAvis.save_html(preset.visualization_data, f)


class LDAStatisticsMarshaler(Marshaler[GensimLDAStage]):
    OUTPUT_NAME: ClassVar[str] = "lda.png"

    _output_manager: OutputManager
    _config: dict

    def __init__(self, output_manager: OutputManager, global_config: dict) -> None:
        self._output_manager = output_manager
        self._config = global_config["topic"]["lda"]

    @override
    def marshal(self, stage: GensimLDAStage) -> None:
        topic_nums = range(self._config["min_topics"], self._config["max_topics"] + 1)
        perplexity_scores = [stat.perplexity for stat in stage.statistics]
        coherence_scores = [stat.coherence for stat in stage.statistics]

        path = self._output_manager.filepath(LDAStatisticsMarshaler.OUTPUT_NAME)
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(self._config["figure_width"], self._config["figure_height"]))
        ax1.plot(topic_nums, perplexity_scores, "bo-")
        ax1.set_title("Perplexity")
        ax1.set_xlabel("Number of Topics")
        ax1.set_ylabel("Perplexity")
        ax1.grid(True)

        ax2.plot(topic_nums, coherence_scores, "ro-")
        ax2.set_title("Coherence")
        ax2.set_xlabel("Number of Topics")
        ax2.set_ylabel("Coherence")
        ax2.grid(True)

        plt.tight_layout()
        plt.savefig(path, dpi=300, bbox_inches="tight")
        plt.close()


class LDAKeywordsExcelMarshaler(Marshaler[GensimLDAStage]):
    """
    Marshaler for LDA keywords.
    """

    _output_dir: pathlib.Path

    def __init__(self, output_manager: OutputManager) -> None:
        self._output_dir = output_manager.filepath("lda-excel")
        if not self._output_dir.exists():
            self._output_dir.mkdir(parents=True)

    @override
    def marshal(self, stage: GensimLDAStage) -> None:
        for current_pass in stage.passes:
            pass_folder = self._output_dir / f"{len(current_pass)}_topics"
            if not pass_folder.exists():
                pass_folder.mkdir(parents=True)
            for topic_index, topic in enumerate(current_pass):
                words = []
                probabilities = []
                for word, probability in topic:
                    words.append(word)
                    probabilities.append(probability)
                dataframe = pd.DataFrame({"Word": words, "Probability": probabilities})
                dataframe.to_csv(pass_folder / f"topic_{topic_index}.csv", index=False)
